#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import glob
import time
import multiprocessing
import torch

sys.path.append('.')

"""
    将子FaceID的人脸图片数量，小于split_face_count的分为一份，大于等于split_face_count分为另一份；
"""


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 100000:
            #     break

            # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            # file_name = lines[2]
            # mode = int(lines[3])

            # feature = np.fromfile(feat_file, dtype=np.float32, count=feat_dim).reshape(1, feat_dim)
            feature = feat_file.read(feat_dim*4)

            subfaceid_file_features = faceid_subfaceid_file_features.get(faceid, None)
            if subfaceid_file_features is None:
                files = [line]
                features = [feature]
                subfaceid_file_features = {subfaceid: [files, features]}
                faceid_subfaceid_file_features[faceid] = subfaceid_file_features
            else:
                file_features = subfaceid_file_features.get(subfaceid, None)
                if file_features is None:
                    files = [line]
                    features = [feature]
                    subfaceid_file_features[subfaceid] = [files, features]
                else:
                    file_features[0].append(line)
                    file_features[1].append(feature)

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def main():
    parser = argparse.ArgumentParser(description="Split Face ReID Image")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/Feature.sel.FR.FR.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data/data/FaceImage/Feature.sel.norm.FR.FR.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage", help="path to output", type=str
    )
    parser.add_argument(
        "--feature_dim", default=256, help="feature dim", type=int
    )
    parser.add_argument(
        "--split_face_count", default=5, help="split face count", type=int
    )
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim)

    out_file_list_path0 = os.path.join(output_dir, 'Feature.FR.split0.txt')
    out_feature_file_path0 = os.path.join(output_dir, 'Feature.FR.split0.dat')
    out_file_list_path1 = os.path.join(output_dir, 'Feature.FR.split1.txt')
    out_feature_file_path1 = os.path.join(output_dir, 'Feature.FR.split1.dat')

    out_file_list0 = open(out_file_list_path0, 'w')
    out_feature_file0 = open(out_feature_file_path0, 'wb')
    out_file_list1 = open(out_file_list_path1, 'w')
    out_feature_file1 = open(out_feature_file_path1, 'wb')

    subface_count0 = 0
    subface_count1 = 0
    subface_image_count0 = 0
    subface_image_count1 = 0
    face_count = 0
    for faceid, subfaceid_file_features in faceid_subfaceid_file_features.items():
        for subfaceid, file_features in subfaceid_file_features.items():
            files, features = file_features
            if args.split_face_count <= len(files):
                subface_count0 += 1
                for file, feature in zip(files, features):
                    out_file_list0.write(file)
                    out_feature_file0.write(feature)
                    subface_image_count0 += 1
            else:
                subface_count1 += 1
                for file, feature in zip(files, features):
                    out_file_list1.write(file)
                    out_feature_file1.write(feature)
                    subface_image_count1 += 1
        face_count += 1
    out_file_list0.close()
    out_feature_file0.close()
    out_file_list1.close()
    out_feature_file1.close()
    os.system('chmod a+wr {}'.format(out_file_list_path0))
    os.system('chmod a+wr {}'.format(out_feature_file_path0))
    os.system('chmod a+wr {}'.format(out_file_list_path1))
    os.system('chmod a+wr {}'.format(out_feature_file_path1))

    print('face_count {}, subface_count0 {}, subface_image_count0 {}, subface_count1 {}, subface_image_count1 {}'.format(face_count, subface_count0, subface_image_count0, subface_count1, subface_image_count1))
    print('finish!')


if __name__ == '__main__':
    main()
